package com.campus.counseling.service.impl;

import com.campus.counseling.config.LSTMConfig;
import com.campus.counseling.service.DL4JModelService;
import com.campus.counseling.service.Word2VecService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.LSTM;
import org.deeplearning4j.nn.conf.layers.RnnOutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.iterator.TestDataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.INDArrayIndex;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.io.File;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

@Slf4j
@Service
@RequiredArgsConstructor
public class DL4JModelServiceImpl implements DL4JModelService {

    private final LSTMConfig lstmConfig;
    private final Word2VecService word2VecService;
    private MultiLayerNetwork model;
    private boolean modelLoaded = false;

    @PostConstruct
    public void init() {
        try {
            log.info("开始初始化LSTM模型...");
            
            // 确保模型目录存在
            File modelDir = new File(lstmConfig.getModelPath()).getParentFile();
            if (!modelDir.exists()) {
                log.info("创建模型目录: {}", modelDir.getAbsolutePath());
                modelDir.mkdirs();
            }
            
            // 尝试加载已有模型
            File modelFile = new File(lstmConfig.getModelPath());
            if (modelFile.exists()) {
                try {
                    log.info("尝试加载已有模型: {}", modelFile.getAbsolutePath());
                    model = ModelSerializer.restoreMultiLayerNetwork(modelFile);
                    modelLoaded = true;
                    log.info("已加载现有模型");
                } catch (Exception e) {
                    log.error("加载已有模型失败，将创建新模型: {}", e.getMessage());
                    model = createModel();
                    saveModel();
                }
            } else {
                log.info("模型文件不存在，创建新模型");
                model = createModel();
                saveModel();
            }
            
            warmupModel();
            log.info("LSTM模型初始化完成");
        } catch (Exception e) {
            log.error("模型初始化失败", e);
            throw new RuntimeException("模型初始化失败: " + e.getMessage());
        }
    }

    private void warmupModel() {
        try {
            log.info("开始模型预热...");
            
            // 使用三种不同情绪的预热样本
            List<String> warmupTexts = Arrays.asList(
                "今天心情特别好，感到非常开心和幸福！",  // 积极情绪
                "最近很压抑，总是感到难过和焦虑。",      // 消极情绪
                "今天去图书馆学习，天气还不错。"         // 中性情绪
            );
            
            log.info("预热样本数量: {}", warmupTexts.size());
            
            // 提取特征并预测
            INDArray features = extractFeatures(warmupTexts);
            log.info("预热特征维度: {}", Arrays.toString(features.shape()));
            
            INDArray predictions = predict(features);
            log.info("预热预测结果: {}", predictions);
            
            log.info("模型预热完成");
            
        } catch (Exception e) {
            log.error("模型预热失败: ", e);
            throw new RuntimeException("模型预热失败: " + e.getMessage());
        }
    }

    @Override
    public INDArray extractFeatures(List<String> texts) {
        try {
            log.info("开始提取特征，文本数量: {}", texts.size());
            
            // 创建特征矩阵 [batchSize, timeSteps=1, inputSize]
            INDArray features = Nd4j.zeros(texts.size(), 1, lstmConfig.getInputSize());  // 改为205维
            
            for (int i = 0; i < texts.size(); i++) {
                String text = texts.get(i);
                
                // 1. 提取Word2Vec特征
                INDArray word2vecFeatures = word2VecService.getSentenceVector(text);
                log.debug("文本: {}, Word2Vec向量维度: {}", text, Arrays.toString(word2vecFeatures.shape()));
                
                // 2. 提取情感特征
                double[] emotionFeatures = extractEmotionFeatures(text);
                
                // 3. 组合特征
                // 复制Word2Vec特征
                features.put(new INDArrayIndex[]{
                    NDArrayIndex.point(i),
                    NDArrayIndex.point(0), 
                    NDArrayIndex.interval(0, lstmConfig.getVectorSize())
                }, word2vecFeatures);
                
                // 添加情感特征
                for (int j = 0; j < emotionFeatures.length; j++) {
                    features.putScalar(new int[]{i, 0, lstmConfig.getVectorSize() + j}, emotionFeatures[j]);
                }
            }
            
            log.info("特征提取完成，shape: {}", Arrays.toString(features.shape()));
            return features;
            
        } catch (Exception e) {
            log.error("特征提取失败: ", e);
            throw new RuntimeException("特征提取失败: " + e.getMessage());
        }
    }

    private double[] extractEmotionFeatures(String text) {
        // 简单的情感词统计
        double[] features = new double[5];  // 固定5个情感特征维度
        
        // 积极情感词
        String[] positiveWords = {"开心", "希望", "充实", "愉快", "好"};
        // 消极情感词
        String[] negativeWords = {"难过", "痛苦", "绝望", "失败", "死"};
        // 程度词
        String[] degreeWords = {"很", "非常", "特别"};
        // 否定词
        String[] negationWords = {"不", "没", "别"};
        
        // 统计词频
        for (String word : positiveWords) {
            if (text.contains(word)) {
                features[0] += 1.0;
            }
        }
        for (String word : negativeWords) {
            if (text.contains(word)) {
                features[1] += 1.0;
            }
        }
        for (String word : degreeWords) {
            if (text.contains(word)) {
                features[2] += 1.0;
            }
        }
        for (String word : negationWords) {
            if (text.contains(word)) {
                features[3] += 1.0;
            }
        }
        
        // 计算文本长度特征
        features[4] = text.length() / 100.0;  // 归一化文本长度
        
        return features;
    }

    @Override
    public INDArray predict(INDArray features) {
        try {
            if (model == null) {
                loadModel();
            }
            
            // 确保输入维度正确
            if (features.shape()[2] != lstmConfig.getInputSize()) {
                throw new IllegalArgumentException(String.format(
                    "输入特征维度错误: expected=%d, actual=%d",
                    lstmConfig.getInputSize(), features.shape()[2]));
            }
            
            // 调整维度顺序为 [batchSize, vectorSize, timeSteps]
            features = features.permute(0, 2, 1);
            
            INDArray output = model.output(features);
            log.debug("预测完成，输出shape: {}", Arrays.toString(output.shape()));
            
            return output;
            
        } catch (Exception e) {
            log.error("预测失败: ", e);
            throw new RuntimeException("预测失败: " + e.getMessage());
        }
    }

    @Override
    public void loadModel() {
        try {
            String modelPath = lstmConfig.getModelPath();
            log.info("开始加载模型，路径: {}", modelPath);
            
            // 尝试加载最佳模型
            String bestModelPath = modelPath.replace(".bin", "_best.bin");
            File bestModelFile = new File(bestModelPath);
            File modelFile = new File(modelPath);
            
            if (bestModelFile.exists()) {
                log.info("发现最佳模型，加载: {}", bestModelPath);
                model = ModelSerializer.restoreMultiLayerNetwork(bestModelFile);
            } else if (modelFile.exists()) {
                // 加载普通模型
                log.info("加载普通模型: {}", modelPath);
                model = ModelSerializer.restoreMultiLayerNetwork(modelFile);
            } else {
                // 如果模型文件不存在，创建新模型
                log.info("模型文件不存在，创建新模型");
                model = createModel();
                // 立即保存模型
                saveModel();
            }
            
            modelLoaded = true;
            log.info("模型加载/创建成功");
            
        } catch (Exception e) {
            log.error("模型加载失败", e);
            throw new RuntimeException("模型加载失败: " + e.getMessage());
        }
    }

    @Override
    public void train(DataSet dataSet) {
        try {
            log.info("开始训练LSTM模型...");
            
            // 1. 确保模型已初始化
            if (model == null) {
                model = createModel();
            }
            
            // 2. 调整数据维度
            INDArray features = dataSet.getFeatures();
            INDArray labels = dataSet.getLabels();
            
            // 记录原始维度
            log.info("原始特征维度: {}", Arrays.toString(features.shape()));
            log.info("原始标签维度: {}", Arrays.toString(labels.shape()));
            
            // 调整特征维度为 [batchSize, inputSize, timeSteps]
            if (features.rank() == 3) {
                features = features.permute(0, 2, 1);
            }
            
            // 确保标签是3维 [batchSize, timeSteps, 1]
            if (labels.rank() != 3) {
                if (labels.rank() == 2) {
                    labels = labels.reshape(labels.size(0), 1, 1);
                } else {
                    throw new IllegalStateException("标签维度不正确: " + Arrays.toString(labels.shape()));
                }
            }
            
            log.info("调整后特征维度: {}", Arrays.toString(features.shape()));
            log.info("调整后标签维度: {}", Arrays.toString(labels.shape()));
            
            // 3. 创建新的DataSet
            DataSet adjustedDataSet = new DataSet(features, labels);
            
            // 4. 创建数据迭代器
            DataSetIterator iterator = new TestDataSetIterator(adjustedDataSet, lstmConfig.getBatchSize());
            
            // 5. 训练参数
            int numEpochs = lstmConfig.getNumEpochs();
            double bestLoss = Double.MAX_VALUE;
            int patience = 10;
            int noImprovement = 0;
            
            // 6. 开始训练
            for (int epoch = 0; epoch < numEpochs; epoch++) {
                double epochLoss = 0;
                int batchCount = 0;
                
                // 6.1 批量训练
                iterator.reset();
                while (iterator.hasNext()) {
                    DataSet batch = iterator.next();
                    model.fit(batch);
                    epochLoss += model.score();
                    batchCount++;
                }
                
                // 6.2 计算平均损失
                epochLoss /= batchCount;
                log.info("Epoch {}/{} - loss: {:.5f}", epoch + 1, numEpochs, epochLoss);
                
                // 6.3 早停检查
                if (epochLoss < bestLoss) {
                    bestLoss = epochLoss;
                    noImprovement = 0;
                    // 保存最佳模型
                    saveModel("_best");
                } else {
                    noImprovement++;
                    if (noImprovement >= patience) {
                        log.info("触发早停: {} 个epoch未改善", patience);
                        break;
                    }
                }
            }
            
            // 7. 保存最终模型
            saveModel();
            
        } catch (Exception e) {
            log.error("LSTM模型训练失败: ", e);
            throw new RuntimeException("LSTM模型训练失败: " + e.getMessage());
        }
    }
    
    private MultiLayerNetwork createModel() {
        try {
            log.info("创建新的LSTM模型，输入维度: {}", lstmConfig.getInputSize());
            
            // 创建新的LSTM模型配置
            MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(lstmConfig.getSeed())
                .updater(new Adam(lstmConfig.getLearningRate()))
                .list()
                .layer(0, new LSTM.Builder()
                    .nIn(lstmConfig.getInputSize())
                    .nOut(lstmConfig.getHiddenSize())
                    .activation(Activation.TANH)
                    .build())
                .layer(1, new RnnOutputLayer.Builder()
                    .nIn(lstmConfig.getHiddenSize())
                    .nOut(1)  // 输出维度为1
                    .activation(Activation.SIGMOID)
                    .lossFunction(LossFunctions.LossFunction.MSE)
                    .build())
                .build();
            
            // 初始化模型
            MultiLayerNetwork newModel = new MultiLayerNetwork(conf);
            newModel.init();
            
            log.info("LSTM模型创建完成，结构: {}", newModel.summary());
            return newModel;
            
        } catch (Exception e) {
            log.error("创建LSTM模型失败: ", e);
            throw new RuntimeException("创建LSTM模型失败: " + e.getMessage());
        }
    }

    private void saveModel(String suffix) {
        try {
            String modelPath = lstmConfig.getModelPath();
            if (suffix != null) {
                modelPath = modelPath.replace(".bin", suffix + ".bin");
            }
            log.info("开始保存模型，路径: {}", modelPath);
            
            // 确保目录存在
            File modelFile = new File(modelPath);
            modelFile.getParentFile().mkdirs();
            
            // 保存模型
            ModelSerializer.writeModel(model, modelFile, true);
            
            log.info("模型保存成功: {}", modelPath);
            
        } catch (Exception e) {
            log.error("模型保存失败: ", e);
            throw new RuntimeException("模型保存失败: " + e.getMessage());
        }
    }

    // 重载无参数的方法
    private void saveModel() {
        saveModel(null);
    }

} 